# -*- coding: utf-8 -*-
#这个文件包括一些跟分块有关的类。

#mfcc是基本的分块类，包含分块算法的一些基本操作。
import amber,molecule,gauss
import numpy as np
#proRes(i,j)是一个元组，i代表蛋白的起始残基号，j代表末尾残基号（残基号最小为1)
#micro是小分子的始末残基号
def getProANDmicro(proRes,microRes,topName):
  prmtop = amber.prmtop(topName)
  prmtop.read() #读取prmtop中的相应信息。
#  print '原子总数= ',len(coordinate)
  proResName = prmtop.residueName[proRes[0]-1:proRes[1]]
#  print '残基总数 ',len(prmtop.residuePointer)
#  print '蛋白总数 ',proRes[1]
  proResIndex = prmtop.residuePointer[proRes[0]-1:proRes[1]]
#  print '残基切出数 ',len(proResIndex)
  proHead = proResIndex[proRes[0]-1][0]
  proTail = proResIndex[proRes[1]-1][1]
#  print 'proHead= ',proHead
#  print 'proTail= ',proTail
  proAtomName = prmtop.atomName[proHead-1:proTail]
  proAtomElement = prmtop.element[proHead-1:proTail]
  microHead = prmtop.residuePointer[microRes[0]-1][0]
  microTail = prmtop.residuePointer[microRes[1]-1][1]
#  print 'microHead= ',microHead
#  print 'microTail= ',microTail
  microElement = prmtop.element[microHead-1:microTail]
#生成蛋白对象
  protein = molecule.protein(proAtomName,proAtomElement,proResName,proResIndex,None,proHead,proTail)
#生成小分子对象
  micromolecule = molecule.micromolecule(microElement,None,microHead,microTail,'micro')
  return (protein,micromolecule)

def addproANDmicroCrd(protein,micro,crdName):
  coordinate = amber.readAmberInpcrd(crdName)#读取坐标信息。
  coordinate = coordinate.reshape(coordinate.size/3,3)
  protein.coordinate = coordinate[protein.head-1:protein.tail]
  micro.coordinate = coordinate[micro.head-1:micro.tail]

#给残基加帽子
#residues是一个字典，键值对是残基号到残基对象的映射。
#head和tail是两个元组，分别记录每条链的头尾残基序号。
def addHcap(residues,protein,head,tail):
  for m in residues:
    resHead = protein.residueIndex[m-1][0]
    resTail = protein.residueIndex[m-1][1]
    if m in head:
      residues[m].capElement = ('H')
      residues[m].capCoordinate = (protein.coordinate[resTail])
    elif m in tail:
      residues[m].capElement = ('H')
      residues[m].capCoordinate = (protein.coordinate[resHead-2])
    else:
      residues[m].capElement = ('H','H')
      residues[m].capCoordinate = (protein.coordinate[resHead-2],protein.coordinate[resTail])
  

#读取amber中输出的力
def readForce(name):                                                             
  strings = open(name,'r')
  string = strings.readlines()
  strings.close()                                            
  string = [line[14:] for line in string[1:]]                                    
  print string[:10]                                                              
  string = ''.join(string)                                                       
  force = np.fromstring(string,sep=' ')                                              
  force = force.reshape(force.size/3,3)                                              
  return force 

def writeForce(name,force,energy):                                               
  string = open(name,'w')                                                        
  string.write('Total energy is:\n')                                             
  string.write('{0:16.8f}\n'.format(energy))                                     
  string.write('Forces\n')                                                       
  for m in xrange(len(force)):                                                   
    string.write('{0:>10}{1:20.10f}{2:20.10f}{3:20.10f}\n'.format(m+1,force[m][0],force[m][1],force[m][2]))                                                                                                                            
  string.close() 

def calcSingle(fragments):
    single = {}
    for m in fragments:
      print 'calc residue =',m
      fragmentGauss = gauss.gauss()
      fragment = fragments[m]
      if hasattr(fragment,'capElement'):
        fragment0 = (fragment.coordinate,fragment.elementName)
        fragment1 = (fragment.capCoordinate,fragment.capElement)
        fragment = (fragment0,fragment1)
      else:
        fragment = ((fragment.coordinate,fragment.elementName),)
      fragmentName = 'res'+str(m)
      fragmentGauss.setBasic(fragmentName,1,0,fragment)
      fragmentGauss.setKey()
      fragmentGauss.write()
      fragmentGauss.run()
      fragmentGauss.read()
      single[m] = fragmentGauss
      print fragmentGauss.calcForce*1186.219
    return single

def calcDouble(fragments):
  double = {}
  micro = fragments.pop('micro')
  fragMicro = (micro.coordinate,micro.elementName)
  for m in fragments:
    print 'double= ',m
    fragmentGauss = gauss.gauss()
    fragment = fragments[m]
    fragment0 = (fragment.coordinate,fragment.elementName)
    if hasattr(fragment,'capElement'):
      fragment1 = (fragment.capCoordinate,fragment.capElement)
      fragment = (fragment0,fragment1,fragMicro)
    else:
      fragment = (fragment0,fragMicro)
    fragmentName = 'res'+str(m)+'micro'
    fragmentGauss.setBasic(fragmentName,1,0,fragment)
    fragmentGauss.setKey()
    fragmentGauss.write()
    fragmentGauss.run()
    fragmentGauss.read()
    double[m] = fragmentGauss
    print fragmentGauss.calcForce*1186.219
  fragments['micro'] = micro
  return double

def calcFull(fragments):
  micro = fragments.pop('micro')
  fragment = []
  print 'fullForce='
  for m in fragments:
    fragment0 = (fragments[m].coordinate,fragments[m].elementName)
    fragment.append(fragment0)
    if hasattr(fragments[m],'capElement'):
      fragment1 = (fragments[m].capCoordinate,fragments[m].elementName)
      fragment.append(fragment1)
  fragment.append((micro.coordinate,micro.elementName))
  fragmentGauss = gauss.gauss()
  fragmentName = 'full'
  fragmentGauss.setBasic(fragmentName,1,0,fragment)
  fragmentGauss.setKey()
  fragmentGauss.write()
  fragmentGauss.run()
  fragmentGauss.read()
  fullForce = fragmentGauss.calcForce[-3:]*1186.219
  print fullForce
  return fullForce


def sumResult(single,double,fragments):
  microSingleForce = single['micro'].calcForce
  microDoubleForces = {}
  for m in double:
    forces = double[m].calcForce
    fragment = fragments[m]
    if hasattr(fragment,'capElement'):
      length = len(fragment.elementName)+len(fragment.capElement)
    else:
      length = len(fragment.elementName)
    microDoubleForce = forces[length:]
    microDoubleForces[m] = microDoubleForce-microSingleForce
  
  for m in microDoubleForces:
    microSingleForce += microDoubleForces[m]
  return microSingleForce
  
def run(proRes,microRes,topName,crdName):
  pro,micro = getProANDmicro(proRes,microRes,topName)
  addproANDmicroCrd(pro,micro,crdName)
#最临近的残基
  close = micro.closestResidue(pro)
#  close = [1,4]
  print 'close=',close
  print 'zzz'
#  closeRes = {}
#  for m in close:
#    closeRes[m] = pro.cutResidue(m)
##给残基加上帽子
##  capRes = addHcap(closeRes,pro,(1),(220))
#  capRes = closeRes
##加入小分子对象
#  capRes['micro'] = micro
##  force = readForce('MM_force.dat')
##  energy = 0
#  single = calcSingle(capRes)
#  double = calcDouble(capRes)
#  full = calcFull(capRes)
#  microForce = sumResult(single,double,capRes)
#  mfccForce = microForce*1186.219
#  print 'mfccForce='
#  print mfccForce
#  print '全体-分块='
#  delta = mfccForce-full
#  delta = np.sum(delta**2,axis=1)
#  delta = np.sqrt(delta)
#  print delta
#  
#  force[micro.head-1:micro.tail] = microForce
#  writeForce('mix_force.dat',force,100)
    
  
  
  
if __name__ == '__main__':
  def testGetProANDmicro():
    pro,micro = getProANDmicro((1,11701),(11702,11702),'1emv1.prmtop')
    addproANDmicroCrd(pro,micro,'1emv1.crd')
    print pro.residueName[:10]
    print micro.elementName
    print pro.coordinate[:10]
    print pro.coordinate[-10:]
    print micro.coordinate

#  run((1,5),(6,6),'waterbox.prmtop','waterbox.crd')
  run((1,11701),(11702,11702),'1emv1.prmtop','1emv1.inpcrd')
